package learning.crf.training;

import learning.crf.inference.IParseScorer;
import learning.data.Dataset;
import learning.data.document.SequenceDocument;
import learning.util.SparseVector;

public class BasicParseScorer implements IParseScorer {

	public static enum Type { ALL_TERMINALS, TRUE_TERMINALS }; 

	static boolean learnTransitionFeatures = true; //false;
	private int[][] transitionFeatures;
	private Dataset<SequenceDocument> dataset;
	
	private SequenceDocument doc;
	private CRFParameters params;
	private Type type;
	
	public BasicParseScorer(Type type) {
		this.type = type;
	}

	public BasicParseScorer(Type type, boolean learnTransitionFeatures) {
		this.type = type;
		this.learnTransitionFeatures = learnTransitionFeatures;
	}

	// transition + output scores
	public float scoreTransition(int pos, int previousState, int state) {

		if (type == Type.ALL_TERMINALS) {
			float sum = 0;
	        SparseVector parameters = params.transitionParameters[previousState][state];
	        
	        int prev = -1;
	        
	        if (learnTransitionFeatures) {
	        	// hack:
	        	// update the last feature to reflect transition
	        	
	        	int last = doc.features[pos].ids.length-1;
	        	prev = doc.features[pos].ids[last];
	        	
	        	//System.out.println("was: " + dataset.getFeatureName(doc.features[pos].ids[last]) + " " + doc.features[pos].vals[last]);
	        	doc.features[pos].ids[last] = transitionFeatures[previousState][state];
		        sum += parameters.dotProduct(doc.features[pos]);
		        doc.features[pos].ids[last] = prev;

	        } else {
		        sum += parameters.dotProduct(doc.features[pos]);
	        	
	        }
	        
	        return sum;
		} else if (type == Type.TRUE_TERMINALS) {
			
			// not implemented yet
			
			return Float.NEGATIVE_INFINITY;
		}
		return Float.NaN;
	}
	
	public SparseVector getFeatures(int pos, int previousState, int state) {
		// vector is only valid until next call of this method
        if (learnTransitionFeatures) {
        	int last = doc.features[pos].ids.length-1;
        	doc.features[pos].ids[last] = transitionFeatures[previousState][state];
        }
		return doc.features[pos];
	}
	

	public void setDataset(Dataset<SequenceDocument> dataset) {
		if (learnTransitionFeatures) {
			this.dataset = dataset;
			
			int d =  dataset.numLabels() + 1;		
			transitionFeatures = new int[d][d];

			for (int i=0; i < d; i++)
				for (int j=0; j < d; j++) {
					//System.out.println("checking for: " + "_transition_" + i + "_" + j);
					transitionFeatures[i][j] = dataset.getFeatureId("_transition_" + i + "_" + j);
				}
		}
	}
	
	public void setDocument(SequenceDocument doc) {
		this.doc = doc;
	}

	public void setParameters(CRFParameters params) {
		this.params = params;
	}	
}
